#!/usr/bin/python3
# coding=utf-8
import openai
from args import *
from flask import request, jsonify
from flask import Flask
import json
from prompts import *
from flask_cors import cross_origin

client = openai.Client(
    api_key="cannot be empty",
    base_url=f"http://{xinference_url}:{xinference_port}/v1"
)


def instruct_model(prompt, model_uid=model_entities_uid, max_tokens=512, temperature=0.7):
    response =  client.completions.create(
        model=model_uid,  # 可以根据需求选择不同的引擎
        prompt=prompt,
        max_tokens=max_tokens,  # 设置响应的最大 token 数
        n=1,  # 返回一个响应
        stop=None,  # 不设置特定的停止序列
        temperature=temperature,  # 控制生成文本的随机性
    )
    return response.choices[0].text.strip()


app = Flask(__name__)
app.debug = False


@app.route('/model/extract/location', methods=['post'])
@cross_origin()
def post_http_entities():
    req = request.get_json()
    context = req.get('context', '')
    prompt = get_location_prompt(context)
    response = instruct_model(prompt)
    try:
        data = json.loads(response)
    except:
        return {
            "code": 200,
            "data": {"person": [], "organization": [], "place": [], "country": []},
            "message": "解析json错误"
        }
    ret = {
        "code": 200,
        "data": data,
        "message": "success"
    }
    return jsonify(ret)


def check_lis_string(lis: list[str], string: str):
    for l in lis:
        if l in string:
            return [l]
    return ["其他"]


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=port)
    # 这里指定了地址和端口号。也可以不指定地址填0.0.0.0那么就会使用本机地址ip
